import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import os
# import coremltools as ct

from datetime import datetime, timedelta

from torch.utils.data import Dataset, DataLoader


# 数据预处理类
class PeriodDataset(Dataset):
    def __init__(self, data, seq_length=4):
        self.data = data
        BM1 = self.data['is_period'] == 1
        self.data['diff_days'] = self.data.apply(lambda row: self.get_diff_days(row), axis=1)
        max_diff_days_info = self.data.groupby('start_date').diff_days.max()
        is_period_info = self.data[BM1].groupby('start_date').is_period.count()
        self.diff_dayss = max_diff_days_info.to_numpy()
        self.is_periods = is_period_info.to_numpy()
        print(self.diff_dayss)
        print(self.is_periods)

        self.features = []
        # 构建序列样本
        for i in range(len(self.diff_dayss)):
            self.features.append((float(self.diff_dayss[i]),float(self.is_periods[i])))

        tensor_data = torch.tensor(self.features)
        mean_tensor = torch.mean(tensor_data, dim=0).numpy()
        std_tensor = torch.std(tensor_data, dim=0).numpy()
        print(mean_tensor)
        print(std_tensor)

        self.samples = []
        for i in range(len(self.features) - seq_length):
            x = self.features[i:i + seq_length]
            y = self.features[i + seq_length]
            self.samples.append((x, y))
        # 归一化参数（每个特征单独归一化）

        self.mean = mean_tensor
        self.std = std_tensor
        print(self.features)
        print(self.mean, self.std)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        x, y = self.samples[idx]
        # 对每个特征进行归一化
        x = (x - self.mean) / self.std
        y = (y - self.mean) / self.std
        return (
            torch.FloatTensor(x),
            torch.FloatTensor(y),
        )

    def get_diff_days(self,row):
        return (datetime.strptime(row['date'], '%Y-%m-%d') - datetime.strptime(row['start_date'], '%Y-%m-%d')).days


# LSTM模型定义
class PeriodPredictor(nn.Module):
    def __init__(self, input_size=2,  # 对应2个特征
                 hidden_size=1,
                 num_layers=2,
                 output_size=2):  # 预测2个特征
        super().__init__()
        self.bn = nn.BatchNorm1d(input_size)
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # x shape: (batch_size, seq_length, input_size)
        out, _ = self.lstm(x)
        out = self.fc(out[:,-1])  # 取最后一个时间步
        return out

def plot_time_series(actual_diff_days, actual_is_period,forecast_diff_days,forecast_is_period, title="时间序列预测"):
    """
    绘制时间序列实际值与预测值对比图
    :param actual: 实际值序列 (带时间索引的Series)
    :param forecast: 预测值序列 (带时间索引的Series)
    """
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统可用
    # plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # Linux系统可用 
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    
    plt.figure(figsize=(12, 6))
    
    # 绘制实际值（蓝色实线）
    plt.plot(actual_diff_days.index, actual_diff_days.values, 'b-', label='实际周期天数')
    plt.plot(actual_is_period.index, actual_is_period.values, 'g-', label='实际经期天数')
    
    # 绘制预测值（红色虚线）
    plt.plot(forecast_diff_days.index, forecast_diff_days.values, 'r--', label='预测周期天数')
    plt.plot(forecast_is_period.index, forecast_is_period.values, 'm--', label='预测经期天数')
    
    plt.title(title)
    plt.xlabel('时间序列')
    plt.ylabel('天数')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    
    # 保存图片（可选）
    plt.savefig('time_series_forecast.png')
    plt.show()

if __name__ == "__main__":
    test_df = pd.read_csv('period_full_data.csv')
    input_size = 1
    # 准备数据集
    dataset = PeriodDataset(test_df, seq_length=input_size)
    dataloader = DataLoader(dataset ,batch_size=2, shuffle=True)

    # 初始化模型
    model_path = 'models/period_predictor.pth'
    if os.path.exists(model_path):
        model = torch.load('models/period_predictor.pth',weights_only=False)
    else:
        model = PeriodPredictor(input_size=2)

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练循环
    for epoch in range(1000):
        for x, y in dataloader:
            optimizer.zero_grad()
            outputs = model(x)

            loss = criterion(outputs, y)
            loss.backward()
            optimizer.step()
        if epoch % 10 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

    # 保存整个模型
    torch.save(model, 'models/period_predictor.pth')
    
    # 新增Core ML模型转换
    # mlmodel_path = 'models/PeriodPredictor.mlmodel'
    # try:
    #     # 将模型设置为评估模式
    #     model.eval()
    #
    #     # 创建示例输入（根据实际输入形状调整）
    #     example_input = torch.randn(1, input_size, 2)  # (batch_size, seq_length, input_features)
    #     print(example_input)
    #     # 转换为TorchScript
    #     traced_model = torch.jit.trace(model, example_input)
    #
    #     # 检查输入形状是否匹配模型预期
    #     print(f"模型输入形状要求: {model.lstm.input_size} -> 实际输入特征数应为2")
    #     print(f"LSTM层参数: {dict(model.named_parameters())}")  # 查看模型参数验证结构
    #     print(f"示例输入形状: {example_input.shape}")
    #     # 应该输出类似：示例输入形状: torch.Size([1, 3, 2]) 表示 (batch, seq_len, features)
    #
    #     # 转换为Core ML模型
    #     mlmodel = ct.convert(
    #         traced_model,
    #         inputs=[ct.TensorType(
    #             name="input",
    #             shape=ct.Shape(shape=(
    #                 1,  # 固定batch_size
    #                 ct.RangeDim(lower_bound=1, upper_bound=10, default=1),  # 序列长度范围
    #                 2    # 特征数固定
    #             )),
    #             dtype=np.float32
    #         )],
    #         outputs=[ct.TensorType(name="output", dtype=np.float32)],
    #         convert_to="neuralnetwork",  # 使用旧版格式
    #         compute_units=ct.ComputeUnit.CPU_ONLY
    #     )
    #
    #     # 添加元数据
    #     mlmodel.author = "Your Name"
    #     mlmodel.short_description = "月经周期预测模型"
    #     mlmodel.input_description["input"] = f"归一化后的时序数据 (seq_length={input_size}, features=2)"
    #     mlmodel.output_description["output"] = "预测的周期和经期天数 (2个特征值)"
    #
    #     # 创建保存目录
    #     os.makedirs('models', exist_ok=True)
    #     # 保存Core ML模型
    #     mlmodel.save(mlmodel_path)
    #     print(f"Core ML模型已保存至 {mlmodel_path}")
    # except Exception as e:
    #     print(f"Core ML转换失败: {str(e)}")

    m = dataset.mean
    s = dataset.std
    # 预测示例修改
    with torch.no_grad():
        actual_diff_days_data = []
        actual_is_period_data = []
        forecast_diff_days_data = []
        forecast_is_period_data = []

        for i in range(len(dataset.features) - input_size):
            test_data = dataset.features[i: i + input_size]
            test_input = (np.array(test_data) - dataset.mean) / dataset.std
            test_input = torch.FloatTensor(test_input)
            prediction = model(test_input)
            # 反归一化
            denorm_pred = prediction.numpy() * dataset.std + dataset.mean
            actual_diff_days_data.append(dataset.features[i + input_size][0])
            actual_is_period_data.append(dataset.features[i + input_size][1])
            forecast_diff_days_data.append(denorm_pred[0])
            forecast_is_period_data.append(denorm_pred[1])

            print(f"下一次周期:{denorm_pred[0]}")
            print(f"下一次多久:{denorm_pred[1]}")

        # 创建示例数据
        actual_diff_days = pd.Series(actual_diff_days_data,index=range(len(actual_diff_days_data)))
        actual_is_period = pd.Series(actual_is_period_data,index=range(len(actual_is_period_data)))
        forecast_diff_days = pd.Series(forecast_diff_days_data,index=range(len(forecast_diff_days_data)))
        forecast_is_period = pd.Series(forecast_is_period_data,index=range(len(forecast_is_period_data)))

        plot_time_series(actual_diff_days, actual_is_period,forecast_diff_days,forecast_is_period, "健康指标预测")

